import torch
import numpy as np
import pandas as pd
import json
from pathlib import Path
import os
import sys
import re
import logging
from resources.DKT_model.dkt import DKTForward
# 设置日志
logging.basicConfig(level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 统一硬件配置
DEVICE = os.getenv("DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")

# 配置路径
dkt_model_path = './resources/dkt_model_final.pth'
kp_map_path = './resources/knowledge_point_mapping.csv'
knows_tree_path = './knowledge_graph_重庆数学.json'
question_bank_path = './resources/question_bank_processed.csv'

def split_student_data(student_data):
    # 拆分原始字段
    exercise_ids = student_data["exercise_ids"].split(",")
    is_corrects = student_data["is_corrects"].split(",")
    knowledge_points = student_data["knowledge_points"].split(";")
    begin_times = student_data["begin_times"].split(",")

    # 查找最后一个错误题的索引
    last_wrong_index = -1
    for i in range(len(is_corrects) - 1, -1, -1):
        if is_corrects[i] == "0":
            last_wrong_index = i
            break

    if last_wrong_index == -1:
        raise ValueError("没有发现错题（is_corrects中没有0）")

    # 构造 current_incorrect 部分
    current_incorrect = {
        "exercise_id": exercise_ids[last_wrong_index],
        "knowledge_points": knowledge_points[last_wrong_index],
        "begin_time": begin_times[last_wrong_index]
    }

    # 构造 previous_sequence 部分（排除最后一题）
    previous_sequence = {
        "exercise_ids": ",".join(exercise_ids[:last_wrong_index]),
        "is_corrects": ",".join(is_corrects[:last_wrong_index]),
        "knowledge_points": ";".join(knowledge_points[:last_wrong_index]),
        "begin_times": ",".join(begin_times[:last_wrong_index])
    }

    # 返回结构化数据
    return {
        "student_id": student_data["student_id"],
        "current_incorrect": current_incorrect,
        "previous_sequence": previous_sequence
    }

def map_knowledge_points(knowledge_points, kp_map_path):
    """
    根据知识点名称和映射表生成 mapped_knowledge_points。
    
    Args:
        knowledge_points (str): 逗号分隔的知识点名称，如"无穷小量,等价无穷小"
        kp_map_path (str): 知识点映射表路径
    
    Returns:
        str: 逗号分隔的映射ID，如"9,96"
    """
    try:
        if not knowledge_points:
            return ""
        
        # 加载映射表
        kp_df = pd.read_csv(kp_map_path)
        kp2id = dict(zip(kp_df['knowledge_point'], kp_df['mapped_id']))
        
        # 分割知识点（同一题目内的知识点用逗号分隔）
        kps = knowledge_points.split(',')
        mapped_ids = []
        for kp in kps:
            kp = kp.strip()
            if kp in kp2id:
                mapped_ids.append(str(kp2id[kp]))
            else:
                logger.warning(f"知识点 '{kp}' 未在映射表中找到，跳过")
        
        return ','.join(mapped_ids) if mapped_ids else ""
    except Exception as e:
        logger.error(f"知识点映射失败: {e}", exc_info=True)
        raise

def parse_server_request(request_data):
    try:
        request_data = split_student_data(request_data)
        student_id = request_data.get('student_id', '未知')
        current_incorrect = request_data.get('current_incorrect', {})
        previous_sequence = request_data.get('previous_sequence', {})

        if not current_incorrect or not previous_sequence:
            missing = []
            if not current_incorrect:
                missing.append("'current_incorrect'")
            if not previous_sequence:
                missing.append("'previous_sequence'")
            raise ValueError(f"请求数据缺少必要字段: {', '.join(missing)}")

        # 检查 current_incorrect 字段
        for field in ('exercise_id', 'knowledge_points', 'begin_time'):
            if field not in current_incorrect:
                raise ValueError(f"current_incorrect 缺少字段: {field}")

        # 检查或生成 current_incorrect 的 mapped_knowledge_points
        if 'mapped_knowledge_points' not in current_incorrect or not current_incorrect['mapped_knowledge_points']:
            current_incorrect['mapped_knowledge_points'] = map_knowledge_points(
                current_incorrect['knowledge_points'], kp_map_path
            )
            if not current_incorrect['mapped_knowledge_points']:
                raise ValueError(f"current_incorrect 的知识点 {current_incorrect['knowledge_points']} 映射失败")

        student_data = {
            'student_id': student_id,
            'exercise_ids': previous_sequence.get('exercise_ids', ''),
            'is_corrects': previous_sequence.get('is_corrects', ''),
            'knowledge_points': previous_sequence.get('knowledge_points', ''),
            'mapped_knowledge_points': previous_sequence.get('mapped_knowledge_points', ''),
            'begin_times': previous_sequence.get('begin_times', ''),
        }

        # 检查或生成 previous_sequence 的 mapped_knowledge_points
        if not student_data['mapped_knowledge_points']:
            # 按分号分割知识点组（不同题目）
            kp_groups = student_data['knowledge_points'].split(';') if student_data['knowledge_points'] else []
            mapped_kp_groups = []
            for kp_group in kp_groups:
                # 每个 kp_group 是逗号分隔的知识点（同一题目）
                mapped_ids = map_knowledge_points(kp_group, kp_map_path)
                mapped_kp_groups.append(mapped_ids)
            student_data['mapped_knowledge_points'] = ';'.join(mapped_kp_groups)

        return {
            'student_id': student_id,
            'current_incorrect': current_incorrect,
            'previous_sequence': previous_sequence,
            'student_data': student_data
        }
    except Exception as e:
        logger.error(f"parse_server_request 解析失败: {e}", exc_info=True)
        raise

def encode_student_onehot(student_data, max_step=50, num_questions=157):
    try:
        student_id = student_data['student_id']
        answer_time = student_data['begin_times'].split(',')[0] if student_data['begin_times'] else '未知'
        exercise_ids = student_data['exercise_ids']
        is_corrects = student_data['is_corrects']
        knowledge_points = student_data['knowledge_points']
        mapped_kp_str = student_data['mapped_knowledge_points']

        kp_groups = mapped_kp_str.split(';') if mapped_kp_str else []
        corrects = list(map(int, is_corrects.split(','))) if is_corrects else []

        onehot = np.zeros((max_step, 2 * num_questions), dtype=np.float32)

        if not kp_groups or len(kp_groups) != len(corrects):
            logger.warning(f"知识点组数 {len(kp_groups)} 与答案数 {len(corrects)} 不匹配，返回空编码")
            return onehot, student_id, answer_time, {}

        length = len(corrects)
        start_idx = max(0, length - max_step)
        for i, (kp_group, correct) in enumerate(zip(kp_groups[start_idx:], corrects[start_idx:])):
            if kp_group:
                kp_ids = list(map(int, kp_group.split(','))) if kp_group else []
                for kp_id in kp_ids:
                    if kp_id >= num_questions:
                        logger.warning(f"知识点ID {kp_id} 超出范围，跳过")
                        continue
                    index = kp_id if correct > 0 else kp_id + num_questions
                    onehot[i][index] = 1

        metadata = {
            "student_id": student_id,
            "answer_time": answer_time,
            "exercise_ids": exercise_ids,
            "is_corrects": is_corrects,
            "knowledge_points": knowledge_points,
            "mapped_knowledge_points": mapped_kp_str
        }

        return onehot, student_id, answer_time, metadata
    except Exception as e:
        logger.error(f"生成 one-hot 编码失败: {e}", exc_info=True)
        raise

# DKTInference 类（未修改，简化为占位符）
class DKTInference:
    def __init__(self, model_path, num_questions, hidden_dim=128, layer_num=2, dropout=0.2, cell_type='lstm', device=DEVICE):
        try:
            if not os.path.exists(model_path):
                raise FileNotFoundError(f"DKT 模型路径 {model_path} 不存在")
            self.device = device
            self.model = DKTForward(
                input_dim=2 * num_questions,
                hidden_dim=hidden_dim,
                layer_num=layer_num,
                output_dim=num_questions,
                dropout=dropout,
                device=self.device,
                cell_type=cell_type
            ).to(self.device)

            checkpoint = torch.load(model_path, map_location=self.device)
            if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                state_dict = checkpoint["model_state_dict"]
            else:
                raise ValueError(f"模型文件 {model_path} 格式错误，期望包含 'model_state_dict' 键")

            self.model.load_state_dict(state_dict)
            self.model.eval()
            self.num_questions = num_questions
        except Exception as e:
            logger.error(f"初始化 DKT 模型失败: {e}", exc_info=True)
            raise

    def predict_student(self, student_onehot_seq: np.ndarray):
        try:
            if student_onehot_seq.shape[1] != 2 * self.num_questions:
                raise ValueError(f"输入序列维度应为 [seq_len, {2 * self.num_questions}]，实际为 {student_onehot_seq.shape}")

            student_tensor = torch.tensor(student_onehot_seq[np.newaxis, ...], dtype=torch.float32).to(self.device)

            with torch.no_grad():
                logits, _ = self.model(student_tensor)
            probs = torch.sigmoid(logits[0]).cpu().numpy()[-1]
            return probs
        except Exception as e:
            logger.error(f"DKT 预测失败: {e}", exc_info=True)
            raise

#generate_final_prompt
def generate_final_prompt(parsed_data, dkt_model, knows_tree_path, kp_map_path, question_bank_path, max_step=50, num_questions=157):
    try:
        current_incorrect = parsed_data['current_incorrect']
        student_data = parsed_data['student_data']

        # Load knows_tree
        with open(knows_tree_path, 'r', encoding='utf-8') as f:
            knows_tree = json.load(f)
        # Load kp_map
        kp_df = pd.read_csv(kp_map_path)
        kp2id = dict(zip(kp_df['knowledge_point'], kp_df['mapped_id']))

        # Encode student one-hot
        onehot, student_id, _, _ = encode_student_onehot(student_data, max_step, num_questions)

        # Predict knowledge probabilities
        knowledge_probs = dkt_model.predict_student(onehot)

        # Get current knowledge points
        exercise_id = current_incorrect['exercise_id']
        current_kps = current_incorrect['knowledge_points'].split(',')
        current_kp_ids = list(map(int, current_incorrect['mapped_knowledge_points'].split(',')))

        # Construct curr_str
        curr_str = ""
        for kp, kp_id in zip(current_kps, current_kp_ids):
            if kp_id < num_questions:
                curr_str += f"{kp}:{knowledge_probs[kp_id]:.2f}，"

        # Construct prereq_str
        prereq_str = ""
        found_prereq = False
        for kp in current_kps:
            prereqs = knows_tree.get(kp, {}).get('prerequisites', [])
            for pre in prereqs:
                pre_id = kp2id.get(pre)
                if pre_id is not None and pre_id < num_questions:
                    found_prereq = True
                    prereq_str += f"{pre}:{knowledge_probs[pre_id]:.2f}，"

        # Construct the sentence
        if found_prereq:
            sentence = f"根据学生历史作答情况得到本题知识点掌握概率“{curr_str[:-1]}”和其前置知识点掌握概率“{prereq_str[:-1]}”"
        else:
            sentence = f"根据学生历史作答情况得到本题知识点掌握概率“{curr_str[:-1]}”"

        return sentence
    except Exception as e:
        logger.error(f"生成句子失败: {e}", exc_info=True)
        raise
